import os
import tarfile
import urllib.request
import math
from collections import defaultdict, Counter
import re
# 1. 下载与解压
DATA_URL = "http://qwone.com/~jason/20Newsgroups/20news-bydate.tar.gz"
DATA_PATH = "20news-bydate.tar.gz"
EXTRACT_DIR = "data"
def load_data(path):
    texts, labels = [], []
    classes = [d for d in os.listdir(path) if os.path.isdir(os.path.join(path, d))]
    for c in classes:
        folder = os.path.join(path, c)
        files = os.listdir(folder)
        for fname in files:
            with open(os.path.join(folder, fname), 'r', errors='ignore') as f:
                text = f.read()
                texts.append(text)
                labels.append(c)
    print(f"加载 {path}，共 {len(texts)} 条样本，类别数：{len(classes)}")
    return texts, labels
def tokenize(text):
    words = re.findall(r'\b[a-zA-Z]+\b', text.lower())
    return [w for w in words if w not in STOPWORDS]
# 4. 朴素贝叶斯类
class NaiveBayes:
    def __init__(self, alpha=1.0):
        self.alpha = alpha
        self.class_word_counts = defaultdict(Counter)
        self.class_doc_counts = Counter()
        self.vocab = set()

    def fit(self, X, y):
        for text, label in zip(X, y):
            words = tokenize(text)
            self.class_doc_counts[label] += 1
            self.class_word_counts[label].update(words)
            self.vocab.update(words)
        self.total_docs = sum(self.class_doc_counts.values())
        print("训练完成，词汇表大小：", len(self.vocab))

    def predict(self, text):
        words = tokenize(text)
        scores = {}
        V = len(self.vocab)
        for c in self.class_doc_counts:
            log_prob = math.log(self.class_doc_counts[c] / self.total_docs)
            total_words = sum(self.class_word_counts[c].values())
            for w in words:
                count = self.class_word_counts[c][w]
                log_prob += math.log((count + self.alpha) / (total_words + self.alpha*V))
            scores[c] = log_prob
        return max(scores, key=scores.get)
    def score(self, X, y):
        correct = 0
        for text, label in zip(X, y):
            pred = self.predict(text)
            if pred == label:
                correct += 1
        acc = correct / len(y)
        print(f"准确率：{acc:.3f}")
        return acc
# 5. 主程序
if __name__ == "__main__":
    download_dataset()
    extract_dataset()

    train_dir = os.path.join(EXTRACT_DIR, "20news-bydate-train")
    test_dir = os.path.join(EXTRACT_DIR, "20news-bydate-test")

    X_train, y_train = load_data(train_dir)
    X_test, y_test = load_data(test_dir)

    nb = NaiveBayes(alpha=1.0)
    nb.fit(X_train, y_train)
    nb.score(X_test, y_test)
